from datetime import timedelta
from pathlib import Path
from time import time_ns

import click

from storage.blobstore.compression import CompressMethod
from storage.blobstore.crypto import EncryptionAlgorithm
from storage.blobstore.fileblob import FileBlobStore
from storage.blobstore.plyvelblob import PlyvelBlobStore
from storage.compress import compress_file, decompress_file
from storage.ddp import dedup_file, restore_file
from storage.encryption import encrypt_file, decrypt_file
from storage.kvstore.lmdb import LMDBKVStore
from storage.kvstore.plyvel import PlyvelKVStore

_kv_choices = click.Choice([
    'lmdb', 'plyvel', 'none'
])
_blob_choices = click.Choice([
    'file', 'plyvel',
])

_arg_existed_file = click.Path(
    exists=True,
    file_okay=True,
    dir_okay=False,
    readable=True,
    writable=False,
    path_type=Path,
)

_arg_existed_dir = click.Path(
    exists=True,
    file_okay=False,
    dir_okay=True,
    readable=True,
    writable=True,
    path_type=Path,
)


@click.group
def main():
    pass


@main.command
@click.option('-d', '--dry', is_flag=True, help='测试模式，不进行文件写入')
@click.option('-b', '--batch', type=int, default=500, help='批量写批大小')
@click.option('--progress', is_flag=True, help='显示进度条')
@click.option('--block-size', type=int, default=64 * 1024, help='分块大小')
@click.option('--encryption', is_flag=True, help='对文件进行加密')
@click.option('--secret', type=str, help='对文件进行加密的密钥，AES128要求16字节，AES256要求32字节')
@click.option('--compression', is_flag=True, help='对文件进行压缩')
@click.option('--blob', type=_blob_choices, default='file', help='blob存储')
@click.option('--kv', type=_kv_choices, default='plyvel', help='键值对存储')
@click.argument('src', type=_arg_existed_file)
@click.argument('dest', type=_arg_existed_dir)
def dedup(src: Path, dest: Path, block_size: int, dry: bool, progress: bool, batch: int,
          kv: str, blob: str, encryption: bool, secret: str, compression: bool):
    start = time_ns()
    if kv == 'lmdb':
        kv_store = LMDBKVStore(dest / '.lmdb-kv', sync=False)
    elif kv == 'plyvel':
        kv_store = PlyvelKVStore(dest / '.plyvel-kv')
    elif kv == 'none':
        kv_store = None
    else:
        raise Exception(f'unknown kv store type {kv}')

    if blob == 'file':
        if kv_store is None:
            raise Exception('kv store required but not specified')

        blob_store = FileBlobStore(dest / '.file-blob', kv_store=kv_store,
                                   secret=b'',
                                   encrypt=EncryptionAlgorithm.Plaintext,
                                   compress=CompressMethod.Plaintext, batch=batch,
                                   mutex_type=None)
    elif blob == 'plyvel':
        blob_store = PlyvelBlobStore(dest / '.plyvel-blob',
                                     encrypt=EncryptionAlgorithm.Plaintext,
                                     compress=CompressMethod.Plaintext)
    else:
        raise Exception(f'unknown blob store type {blob}')

    dedup_file(src, dest, block_size, blob_store, show_progress=progress, dry=dry)
    elapsed = timedelta(milliseconds=(time_ns() - start) / 1e6)
    print(f'all done! took {elapsed}')
    blob_store.close()
    kv_store.close()


@main.command
@click.option('-d', '--dry', is_flag=True, help='测试模式，不进行文件写入')
@click.option('--progress', is_flag=True, help='显示进度条')
@click.option('--blob', type=_blob_choices, default='file', help='blob存储')
@click.option('--kv', type=_kv_choices, default='plyvel', help='键值对存储')
@click.argument('src', type=_arg_existed_file)
@click.argument('dest', type=_arg_existed_dir)
def restore(src: Path, dest: Path, kv: str, blob: str, dry: bool, progress: bool):
    start = time_ns()
    if kv == 'lmdb':
        kv_store = LMDBKVStore(src.parent / '.lmdb-kv', sync=False)
    elif kv == 'plyvel':
        kv_store = PlyvelKVStore(src.parent / '.plyvel-kv')
    elif kv == 'none':
        kv_store = None
    else:
        raise Exception(f'unknown kv store type {kv}')

    if blob == 'file':
        if kv_store is None:
            raise Exception('kv store required but not specified')

        blob_store = FileBlobStore(src.parent / '.file-blob', kv_store=kv_store, encrypt=EncryptionAlgorithm.Plaintext,
                                   compress=CompressMethod.Plaintext, mutex_type=None)
    elif blob == 'plyvel':
        blob_store = PlyvelBlobStore(src.parent / '.plyvel-blob', encrypt=EncryptionAlgorithm.Plaintext,
                                     compress=CompressMethod.Plaintext)
    else:
        raise Exception(f'unknown blob store type {blob}')

    restore_file(src, dest, blob_store, dry=dry, show_progress=progress)
    elapsed = timedelta(milliseconds=(time_ns() - start) / 1e6)
    print(f'all done! took {elapsed}')


@main.command
@click.option('-p', '--parallel', type=int, default=4, help='并行执行加密的线程数量')
@click.option('-d', '--dry', is_flag=True, help='测试模式，不进行文件写入')
@click.option('--progress', is_flag=True, help='显示进度条')
@click.option('--block-size', type=int, default=128 * 1024, help='分块大小')
@click.argument('src', type=_arg_existed_file)
@click.argument('dest', type=_arg_existed_dir)
def encrypt(src: Path, dest: Path, block_size: int, parallel: int, dry: bool, progress: bool):
    start = time_ns()
    encrypt_file(src, Path(str(dest / src.name) + '.enc'), block_size, EncryptionAlgorithm.AES128CTR,
                 b'0123456701234567',
                 parallel=parallel, show_progress=progress, discard_output=dry)
    elapsed = timedelta(milliseconds=(time_ns() - start) / 1e6)
    print(f'all done! took {elapsed}')


@main.command
@click.option('-p', '--parallel', type=int, default=4, help='并行执行解密的线程数量')
@click.option('-d', '--dry', is_flag=True, help='测试模式，不进行文件写入')
@click.option('--progress', is_flag=True, help='显示进度条')
@click.argument('src', type=_arg_existed_file)
@click.argument('dest', type=_arg_existed_dir)
def decrypt(src: Path, dest: Path, parallel: int, dry: bool, progress: bool):
    start = time_ns()
    decrypt_file(src, Path(str(dest / src.name) + '.dec'), b'0123456701234567',
                 parallel=parallel, show_progress=progress, discard_output=dry)
    elapsed = timedelta(milliseconds=(time_ns() - start) / 1e6)
    print(f'all done! took {elapsed}')


@main.command
@click.option('-s', '--split', type=int, default=4, help='拆分、并行执行压缩的线程数量')
@click.option('-d', '--dry', is_flag=True, help='测试模式，不进行文件写入')
@click.option('--progress', is_flag=True, help='显示进度条')
@click.argument('src', type=_arg_existed_file)
@click.argument('dest', type=_arg_existed_dir)
def compress(src: Path, dest: Path, split: int, progress: bool, dry: bool):
    start = time_ns()
    compress_file(src, dest, split=split, show_progress=progress, discard_output=dry)
    elapsed = timedelta(milliseconds=(time_ns() - start) / 1e6)
    print(f'all done! took {elapsed}')


@main.command
@click.option('-d', '--dry', is_flag=True, help='测试模式，不进行文件写入')
@click.option('--progress', is_flag=True, help='显示进度条')
@click.argument('src', type=_arg_existed_file)
@click.argument('dest', type=_arg_existed_dir)
def decompress(src: Path, dest: Path, dry: bool, progress: bool):
    start = time_ns()
    decompress_file(src, dest, show_progress=progress, dry=dry)
    elapsed = timedelta(milliseconds=(time_ns() - start) / 1e6)
    print(f'all done! took {elapsed}')


if __name__ == '__main__':
    main()
